import torch as pt
import numpy as np
from python_ai.common.xcommon import *

sep('data')
DIM1, DIM2, DIM3 = 2, 3, 4
N = DIM1 * DIM2 * DIM3
data = np.arange(N).reshape(DIM1, DIM2, DIM3).astype(np.float32)
x = pt.Tensor(data)
print(x)

sep('softmax dim=0')
h = pt.nn.Softmax(dim=0)(x)
print(h)
print(h[0] + h[1])
sum = pt.exp(x).sum()
sum01 = pt.exp(x[0]).sum()
sum02 = pt.exp(x[1]).sum()
print(sum01 / sum, sum02 / sum)

sep('softmax dim=1')
h = pt.nn.Softmax(dim=1)(x)
print(h)
print(h[:, 0, :] + h[:, 1, :] + h[:, 2, :])
sum01 = pt.exp(x[:, 0]).sum()
sum02 = pt.exp(x[:, 1]).sum()
sum03 = pt.exp(x[:, 2]).sum()
print(sum01 / sum, sum02 / sum, sum03 / sum)

sep('softmax dim=2')
h = pt.nn.Softmax(dim=2)(x)
print(h)
print(h.sum(dim=2))
sum01 = pt.exp(x[:, :, 0]).sum()
sum02 = pt.exp(x[:, :, 1]).sum()
sum03 = pt.exp(x[:, :, 2]).sum()
sum04 = pt.exp(x[:, :, 3]).sum()
print(sum01 / sum, sum02 / sum, sum03 / sum, sum04 / sum)
